package main

import (
	"database/sql"
	"encoding/json"
	"html/template"
	"log"
	"net/http"
	"net/url"
	"time"

	"github.com/gorilla/sessions"
	"github.com/gorilla/websocket"
	_ "github.com/mattn/go-sqlite3"
)

var (
	db    *sql.DB
	store = sessions.NewCookieStore([]byte("secret-key")) // 使用安全的密钥
)

// Configure upgrader for WebSocket connections.
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func main() {
	// Initialize SQLite database.
	var err error
	store.Options = &sessions.Options{
		Path:     "/",
		MaxAge:   86400, // 1 天有效期
		HttpOnly: true,
	}
	db, err = sql.Open("sqlite3", "../users.db")
	if err != nil {
		log.Fatal("Error opening database: ", err)
	}
	defer db.Close()

	// 添加清理旧消息的代码
	cleanOldMessages()

	// Create users table if it doesn't exist.
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS users (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		username TEXT UNIQUE NOT NULL,
		password TEXT NOT NULL
	)`)
	if err != nil {
		log.Fatal("Error creating table: ", err)
	}

	// 创建好友关系表
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS friendships (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		owner TEXT NOT NULL,
		friend TEXT NOT NULL,
		UNIQUE(owner, friend)
	)`)
	if err != nil {
		log.Fatal("Error creating friendships table: ", err)
	}

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS messages (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		from_user TEXT NOT NULL,
		to_user TEXT NOT NULL,
		message TEXT NOT NULL,
		timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
		is_read INTEGER DEFAULT 0
	)`)
	if err != nil {
		log.Fatal("Error creating messages table: ", err)
	}

	// Route for login (GET/POST) and logout.
	http.HandleFunc("/login", loginHandler)
	http.HandleFunc("/logout", logoutHandler)

	// Serve static assets (login.css, login.js, etc.) without authentication.
	http.Handle("/css/", http.StripPrefix("/css/", http.FileServer(http.Dir("../frontend/css"))))
	http.Handle("/js/", http.StripPrefix("/js/", http.FileServer(http.Dir("../frontend/js"))))

	// Protect chat page, websocket, and addFriend endpoints.
	http.Handle("/", authMiddleware(http.HandlerFunc(indexHandler)))
	http.Handle("/ws", authMiddleware(http.HandlerFunc(handleWebSocket)))
	http.Handle("/addFriend", authMiddleware(http.HandlerFunc(addFriendHandler)))
	http.Handle("/getMessages", authMiddleware(http.HandlerFunc(getMessagesHandler)))
	http.Handle("/markAsRead", authMiddleware(http.HandlerFunc(markAsReadHandler)))

	log.Println("Server started on :8080")
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

// 添加清理旧消息的函数
func cleanOldMessages() {
	// 清理6个月前的未读消息
	sixMonthsAgo := time.Now().AddDate(0, -6, 0)
	result, err := db.Exec(`
        DELETE FROM messages 
        WHERE is_read = 0 
        AND timestamp < ?
    `, sixMonthsAgo)
	if err != nil {
		log.Printf("Error cleaning unread messages: %v", err)
	} else {
		count, _ := result.RowsAffected()
		log.Printf("Cleaned %d unread messages older than 6 months", count)
	}

	// 清理1个月前的已读消息
	oneMonthAgo := time.Now().AddDate(0, -1, 0)
	result, err = db.Exec(`
        DELETE FROM messages 
        WHERE is_read = 1 
        AND timestamp < ?
    `, oneMonthAgo)
	if err != nil {
		log.Printf("Error cleaning read messages: %v", err)
	} else {
		count, _ := result.RowsAffected()
		log.Printf("Cleaned %d read messages older than 1 month", count)
	}
}

// loginHandler handles GET (display login page) and POST (login or register) requests.
func loginHandler(w http.ResponseWriter, r *http.Request) {
	session, _ := store.Get(r, "session")

	// If already logged in, redirect to index.
	if r.Method == "GET" && session.Values["user"] != nil {
		http.Redirect(w, r, "/", http.StatusFound)
		return
	}

	if r.Method == "GET" {
		// Obtain the message from query parameters.
		message := r.URL.Query().Get("message")
		tmpl, err := template.ParseFiles("../frontend/login.html")
		if err != nil {
			http.Error(w, "模板加载失败", http.StatusInternalServerError)
			return
		}
		// Render template with the message data.
		tmpl.Execute(w, struct {
			Message string
		}{Message: message})
		return
	}

	// POST: process login/registration.
	username := r.FormValue("username")
	password := r.FormValue("password")
	if username == "" || password == "" {
		http.Redirect(w, r, "/login?message="+url.QueryEscape("用户名和密码不能为空"), http.StatusFound)
		return
	}
	log.Printf("login: username=%s, password=%s", username, password)

	// Check if user exists.
	row := db.QueryRow("SELECT password FROM users WHERE username = ?", username)
	var existingPassword string
	err := row.Scan(&existingPassword)
	if err != nil {
		if err == sql.ErrNoRows {
			// Register new user.
			_, err = db.Exec("INSERT INTO users(username, password) VALUES(?, ?)", username, password)
			if err != nil {
				http.Redirect(w, r, "/login?message="+url.QueryEscape("注册失败"), http.StatusFound)
				return
			}
		} else {
			http.Redirect(w, r, "/login?message="+url.QueryEscape("数据库错误"), http.StatusFound)
			return
		}
	} else {
		// User exists. Validate password.
		if existingPassword != password {
			http.Redirect(w, r, "/login?message="+url.QueryEscape("密码错误"), http.StatusFound)
			return
		}
	}
	// Set session and redirect to chat page.
	session.Values["user"] = username
	if err := session.Save(r, w); err != nil {
		if r.Header.Get("X-Requested-With") == "XMLHttpRequest" {
			http.Error(w, "无法保存session", http.StatusInternalServerError)
			return
		}
		http.Redirect(w, r, "/login?message="+url.QueryEscape("无法保存session"), http.StatusFound)
		return
	}

	if r.Header.Get("X-Requested-With") == "XMLHttpRequest" {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"redirect":"/"}`))
	} else {
		http.Redirect(w, r, "/", http.StatusFound)
	}
}

// logoutHandler clears the session.
func logoutHandler(w http.ResponseWriter, r *http.Request) {
	session, _ := store.Get(r, "session")
	delete(session.Values, "user")
	session.Save(r, w)
	http.Redirect(w, r, "/login", http.StatusFound)
}

// authMiddleware protects handlers from unauthenticated access.
func authMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		session, _ := store.Get(r, "session")
		if session.Values["user"] == nil {
			log.Printf("Unauthenticated access to %s", r.URL.Path)
			http.Redirect(w, r, "/login", http.StatusFound)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// indexHandler serves the protected chat page.
func indexHandler(w http.ResponseWriter, r *http.Request) {
	session, _ := store.Get(r, "session")
	currentUser, ok := session.Values["user"].(string)
	if !ok {
		http.Redirect(w, r, "/login", http.StatusFound)
		return
	}

	// 查询好友列表
	rows, err := db.Query("SELECT friend FROM friendships WHERE owner = ?", currentUser)
	if err != nil {
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}
	defer rows.Close()

	var friends []string
	for rows.Next() {
		var friend string
		if err := rows.Scan(&friend); err != nil {
			continue
		}
		friends = append(friends, friend)
	}

	// 从请求的查询参数中获取聊天对象
	chatPartner := r.URL.Query().Get("friend")

	tmpl, err := template.ParseFiles("../frontend/index.html")
	if err != nil {
		http.Error(w, "Template parse error", http.StatusInternalServerError)
		return
	}

	data := struct {
		Friends     []string
		ChatPartner string
	}{
		Friends:     friends,
		ChatPartner: chatPartner,
	}
	tmpl.Execute(w, data)
}

type ChatMessage struct {
	From    string `json:"from"`
	To      string `json:"to"`
	Message string `json:"message"`
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Error while upgrading connection:", err)
		return
	}
	defer conn.Close()

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			// 如果连接正常关闭，不记录错误
			if websocket.IsCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				return
			}
			log.Println("Error while reading message:", err)
			break
		}
		// 解析收到的消息（假设客户端发送 JSON 格式数据）
		var chatMsg ChatMessage
		err = json.Unmarshal(msg, &chatMsg)
		if err != nil {
			log.Println("Error decoding JSON message:", err)
			continue
		}
		log.Printf("Received message: %+v", chatMsg)
		// 将 from_user, to_user, message 写入数据库
		_, err = db.Exec("INSERT INTO messages (from_user, to_user, message) VALUES (?, ?, ?)",
			chatMsg.From, chatMsg.To, chatMsg.Message)
		if err != nil {
			log.Println("Error inserting message into DB:", err)
			continue
		}
	}
}

// addFriendHandler handles friend verification requests.
func addFriendHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 解析请求body中的JSON数据
	var req struct {
		Friend string `json:"friend"`
	}
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil || req.Friend == "" {
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	// 查询数据库是否存在该用户
	row := db.QueryRow("SELECT username FROM users WHERE username = ?", req.Friend)
	var friendUsername string
	err = row.Scan(&friendUsername)
	w.Header().Set("Content-Type", "application/json")
	if err != nil {
		if err == sql.ErrNoRows {
			// 用户不存在
			w.Write([]byte(`{"success": false}`))
			return
		}
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}

	// 用户存在，保存好友关系。同时确保一个好友不会重复添加。
	session, _ := store.Get(r, "session")
	owner, ok := session.Values["user"].(string)
	if !ok || owner == "" {
		// 这种情况应该不会发生，因为 /addFriend 已经用 authMiddleware 保护
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	// 插入好友关系，忽略已经存在的关系
	_, err = db.Exec("INSERT OR IGNORE INTO friendships(owner, friend) VALUES(?, ?)", owner, friendUsername)
	if err != nil {
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}
	w.Write([]byte(`{"success": true}`))
}

func getMessagesHandler(w http.ResponseWriter, r *http.Request) {
	session, _ := store.Get(r, "session")
	currentUser, ok := session.Values["user"].(string)
	if !ok {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	// 获取since参数，如果没有则默认获取24小时内的消息
	since := r.URL.Query().Get("since")
	var sinceTime time.Time
	var err error

	if since != "" {
		sinceTime, err = time.Parse(time.RFC3339, since)
		if err != nil {
			// 如果解析失败，使用24小时前作为默认值
			sinceTime = time.Now().Add(-24 * time.Hour)
		}
	} else {
		sinceTime = time.Now().Add(-24 * time.Hour)
	}
	log.Printf("Getting messages since %v", sinceTime)
	// 查询指定时间之后的新消息和所有未读消息
	rows, err := db.Query(`
        SELECT id, from_user, to_user, message, timestamp, is_read 
        FROM messages 
        WHERE ((to_user = ? OR from_user = ?) 
            AND (timestamp > ? OR is_read = 0))
        ORDER BY timestamp DESC
    `, currentUser, currentUser, sinceTime)

	if err != nil {
		log.Printf("Error querying messages: %v", err)
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}
	defer rows.Close()

	// 构造消息列表
	type Message struct {
		ID        int64     `json:"id"`
		FromUser  string    `json:"from_user"`
		ToUser    string    `json:"to_user"`
		Message   string    `json:"message"`
		Timestamp time.Time `json:"timestamp"`
		IsRead    bool      `json:"is_read"`
	}

	var messages []Message
	for rows.Next() {
		var msg Message
		var isRead int
		err := rows.Scan(&msg.ID, &msg.FromUser, &msg.ToUser, &msg.Message, &msg.Timestamp, &isRead)
		if err != nil {
			log.Printf("Error scanning message: %v", err)
			continue
		}
		msg.IsRead = isRead == 1
		messages = append(messages, msg)
	}

	// 设置响应头
	w.Header().Set("Content-Type", "application/json")
	// 返回消息列表
	if err := json.NewEncoder(w).Encode(messages); err != nil {
		log.Printf("Error encoding messages: %v", err)
		http.Error(w, "JSON encoding error", http.StatusInternalServerError)
		return
	}
}

func markAsReadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 获取当前用户
	session, _ := store.Get(r, "session")
	currentUser, ok := session.Values["user"].(string)
	if !ok {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	// 解析请求体
	var req struct {
		Friend string `json:"friend"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	// 更新消息为已读状态
	_, err := db.Exec(`
        UPDATE messages 
        SET is_read = 1 
        WHERE to_user = ? 
        AND from_user = ? 
        AND is_read = 0
    `, currentUser, req.Friend)

	if err != nil {
		log.Printf("Error marking messages as read: %v", err)
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(map[string]bool{"success": true})
}
